import sys
import os
# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint
import matplotlib.animation as animation
import pendulum_config

# 从pendulum_config模块获取参数
l = pendulum_config.l

def create_animation(k, theta0, omega0, t_max=10):
    """
    创建单摆运动动画
    """
    # 计算轨迹
    t = np.linspace(0, t_max, 500)
    y0 = [theta0, omega0]
    sol = odeint(pendulum_config.pendulum_eq, y0, t, args=(k,))
    theta = sol[:, 0]
    
    # 设置图形
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_xlabel('x (m)', fontsize=12)
    ax.set_ylabel('y (m)', fontsize=12)
    ax.set_title('有阻尼单摆运动动画', fontsize=14)
    
    # 固定点
    fixed_point, = ax.plot(0, 0, 'ko', markersize=10)
    
    # 摆杆
    rod, = ax.plot([], [], 'o-', lw=2, markersize=10)
    
    # 轨迹
    trajectory, = ax.plot([], [], 'r-', lw=1, alpha=0.5)
    
    # 轨迹存储
    x_trace, y_trace = [], []
    
    def animate(i):
        # 计算摆的位置
        x = l * np.sin(theta[i])
        y = -l * np.cos(theta[i])
        
        # 更新摆杆
        rod.set_data([0, x], [0, y])
        
        # 更新轨迹
        x_trace.append(x)
        y_trace.append(y)
        if len(x_trace) > 200:  # 只保留最近的200个点
            x_trace.pop(0)
            y_trace.pop(0)
        trajectory.set_data(x_trace, y_trace)
        
        return rod, trajectory
    
    # 创建动画
    anim = animation.FuncAnimation(fig, animate, frames=len(t), interval=50, blit=True, repeat=True)
    
    # 保存动画
    anim.save('damped_pendulum_animation.gif', writer='pillow', fps=20)
    plt.show()
    
    return anim